import sys
sys.path.append('../')
from board_fast import Board,Game,Player,Point,Move
from tensorflow import keras
import numpy as np
from utility.keras_modal import DenseModel
import random
import shutil
from scipy.stats import binom_test
import os
from pathlib import Path
from utility.hdf5 import HDF5
from sgf_parser.sgf_parser_to_game import sgf_parser
import os



class TrainRobot():
    def __init__(self,net='pd_dense',boardSize=9,rand=True):
        self.model=DenseModel(boardSize=boardSize,model=net)
        '''
        if weights_file and rand=False:
            self.model=keras.model.load_model(weights_file)
        '''
        self.boardSize=boardSize
        self.moves=frozenset()  #{(input,output)}
        self.rand=rand
        #self.errCollections=[]  #记录网络走出的异常落子，这些走法都是非法的，要让网络学会（不包含全局同形）[(input,output)]

    def dp_compile(self):
        self.model.p_d_compile()
    def dp_fit(self,x_train,y_train):
        self.model.fit_all_data(x_train,y_train,checkpoint=True)

    def reset(self):
        self.moves=frozenset()

    def load_weights(self,fileName):
        self.rand=False #使用网络权重预测，需要关闭随机策略
        self.model.model_load_weights(fileName)

    def unload_weights(self,fileName):
        self.model.model_save_weights(fileName)

    def isPolicyLegal(self,point,board,player):
        neighbours=board.get_neighbors(point)
        is_eye=True
        for i in neighbours:
            if board.grid.get(i)==None:
                is_eye=False
                break
            elif board.grid.get(i).player != player:
                is_eye=False
                break
            elif len(board.grid.get(i).liberties)<=1:
                is_eye=False
                break
            else:
                pass
        if is_eye:
            return False
        return True

    def predict(self,player,board,reChoose=False,isRandom=None): #game_input的形状是棋盘形状
        if not reChoose:
            '''
            for i in self.moves:
                self.errCollections.append(i)
            '''
            self.moves=frozenset()
        boardSize=self.boardSize
        #mask=np.ones((self.boardSize,self.boardSize))   #让网络学会已有子的地方不能落子,一开始，所有的棋盘点都可落子
        if isRandom:    #这个判断可以强制是的predict采用随机
            rand=True
        else:
            rand=self.rand
        if not np.random.randint(100):   #1%的几率一定使用随机来预测,加入这个扰动，防止陷入局部最优
            rand=True
        if rand:
            npdata=board.print_board(isprint=False)
            mask=(npdata==0).astype(int)
            for i in self.moves:
                mask[i]=False
            if mask.sum()==0:
                return Move(is_pass=True)
            [x,y]=np.where(mask==True)
            moves_=zip(x,y)
            moves=[i for i in moves_]
            random.shuffle(moves)
            point=None
            for i in moves:
                point_=Point(i[0]+1,i[1]+1)
                self.moves=self.moves|{i}
                if self.isPolicyLegal(point_,board,player):
                    point=point_
                    break
            if point is not None:
                move=Move(point=point)
            else:
                move=Move(is_pass=True)
            return move
        else:    #使用网络
            while True:
                npdata=board.print_board(isprint=False)
                mask=(npdata==0).astype(int)
                for i in self.moves:
                    mask[i]=False
                if mask.sum()==0:
                    return Move(is_pass=True)
                if player == Player.black or player == 1:
                    color=np.ones((1,1)).flatten()
                else:
                    color=-1*np.ones((1,1)).flatten()
                npdata=npdata.flatten()
                npdata=np.hstack((npdata,color)).reshape(1,self.boardSize**2+1).astype(int)
                '''
                pred=self.model.model_predict(npdata).reshape(boardSize,boardSize)
                pred_valid=mask*pred
                move=np.unravel_index(np.argmax(pred_valid, axis=None), pred_valid.shape)
                '''
                pred=self.model.model_predict(npdata)
                mask=mask.flatten()
                pred_valid=mask*pred
                pick=random.choices(range(boardSize*boardSize),weights=pred_valid.flatten())  #网络输出取值按比例，不是固定的
                move=np.unravel_index(pick[0], (boardSize,boardSize))
                self.moves=self.moves|{move}
                point=Point(move[0]+1,move[1]+1)
                if self.isPolicyLegal(point,board,player):
                    return Move(point=point)
                else:
                    continue




class PD_Object():
    #def __init__(self):
        #self.weights_current='current_weights.h5'   #保存最新学到的权重
        #self.weights_old='old_weights.h5'   #保存老版本的权重
    def make_samples(self,rounds,bot1,bot2):
        bot1_win=0
        bot2_win=0
        board = Board(size=9)
        game=Game(board)
        for i in range(rounds):
            bot1.reset()
            bot2.reset()
            board.reset()
            game.reset(board)
            #bot1=TrainRobot(rand=False)
            #50%执黑，50%执白
            if np.random.randint(2)==0:
                result=game.run_train(play_b=bot1,play_w=bot2,isprint=False)
                if result=='GameResult.wWin':
                    bot2_win+=1
                else:
                    bot1_win+=1
            else:
                result=game.run_train(play_b=bot2,play_w=bot1,isprint=False)
                if result=='GameResult.wWin':
                    bot1_win+=1
                else:
                    bot2_win+=1
        return bot1_win,bot2_win   #返回bot1和bot2的胜数以及
        '''
        total=bot1_win+bot2_win
        if binom_test(bot1_win, total, 0.5)<.05 and bot1_win/total>.5:
            return 1
        else:
            return -1

            #keras.backend.clear_session()   #没有这句内存要崩
        '''

def play_against_the_other(pd,bot1,bot2,loops):

    #bot1是训练，bot2是进步后才更新的基准
    bot1_win,bot2_win=pd.make_samples(loops,bot1,bot2)
    return bot1_win,bot2_win

def list_all_file(fPath):   #获取目录下的全部文件
    fileList=os.listdir(fPath)
    file_list=[]
    for i in fileList:
        file_dir=Path(fPath+'/'+i)
        if file_dir.is_file():
            file_list.append(file_dir)
    return file_list

def make_tran_data(games_doc,train_file,file_num=0):   #100个文件也就10秒钟最多了，所以不并行了,file_name用来控制并行时候获取多少个文件就退出，暂不启用
    file_lists=list_all_file(games_doc)
    for i in file_lists:

        #'''
        try:
            one_sgf=sgf_parser(i,h5file=train_file)
            one_sgf.setSamples()
            #one_sgf.showGame()
            one_sgf.closeParser()
        except: #如果文件名有雷同，就会出错,但是跳过没关系
            print("error,filename:",i)
        else:
            None
        #''''
        os.remove(i)    #处理完了就删掉


if __name__ == "__main__":
    pd=PD_Object()
    data_file='./dl_p/train_file.h5'   #存放训练数据
    games_doc='./lr_doc/'   #存放对弈的sgf
    weights_current='./dl_p/current_weights.h5'   #保存最新学到的权重
    weights_old='./dl_p/old_weights.h5'   #保存老版本的权重
    #argv[1],argv[2]指定是不是使用
    if len(sys.argv) > 1:
        bot1_isrand=int(sys.argv[1])
    else:
        bot1_isrand=0   #使用初始随机权重
    if len(sys.argv) > 2:
        bot2_isrand=int(sys.argv[2])
    else:
        bot2_isrand=0
    bot1=TrainRobot(rand=bot1_isrand)   #最新的权重
    bot2=TrainRobot(rand=bot2_isrand)   #上一次有进步的权重
    if len(sys.argv) > 3 and int(sys.argv[3])==1:   #这个用在如果训练中断了，可以继续从中断处开始训练
        bot1.load_weights(weights_current)
        bot2.load_weights(weights_old)
    bot1.dp_compile()
    count=0
    while True:
        count+=1
        print("Train times:",count)
        bot1_win,bot2_win=play_against_the_other(pd,bot1,bot2,50)  #这里调整对弈局数
        total=bot1_win+bot2_win
        if binom_test(bot1_win, total, 0.5)<.05 and bot1_win/total>.5:  #bot1显著优秀
            bot1.unload_weights(weights_old)
            bot2.load_weights(weights_old)
            bot2.rand=0
        '''
        elif binom_test(bot1_win, total, 0.5)<.05 and bot2_win/total>.5:    #bot2显著好
            if bot2.rand==0:
                bot1.load_weights(weights_old)
        '''
        else:
            pass
        make_tran_data(games_doc,data_file)
        games=HDF5(data_file,mode='a')
        x_,y_train=games.get_dl_dset()
        games.closeH5()
        os.remove(data_file)    #训练完成后要删掉，后面继续来新的
        x_train=x_[:,:-1]
        player=x_[:,-2]
        winner=x_[:,-1]
        #只学习好的落子法，这个方式不好，但是可以加速学习，实际训练时需要实施下面这个逆向过滤掉不好的下法
        for  i,y in enumerate(player==winner):  #查找落子方是不是胜利方
            if y==False:    #如果不是，就不要学习这一步
                #y_train[i][y_train[i]==0]=1/(y_train[i].size-1)
                #y_train[i][y_train[i]==1]=0
                y_train[i][y_train[i]==1]=-1   #可以设置成-1,和上面注释掉的两句二选一
        bot1.dp_fit(x_train,y_train)
        bot1.unload_weights(weights_current)

    #make_tran_data(games_doc,data_file) #生成用于训练的数据hdf5
